import PIL.Image
import dlib
import face_recognition_models
import numpy

face_detector = dlib.get_frontal_face_detector()

predictor_68_point_model = face_recognition_models.pose_predictor_model_location()
pose_predictor_68_point = dlib.shape_predictor(predictor_68_point_model)

predictor_5_point_model = face_recognition_models.pose_predictor_five_point_model_location()
pose_predictor_5_point = dlib.shape_predictor(predictor_5_point_model)

cnn_face_detection_model = face_recognition_models.cnn_face_detector_model_location()
cnn_face_detector = dlib.cnn_face_detection_model_v1(cnn_face_detection_model)

face_recognition_model = face_recognition_models.face_recognition_model_location()
face_encoder = dlib.face_recognition_model_v1(face_recognition_model)


# step1 获取面部位置
def _raw_face_locations(img, number_of_times_to_upsample=1, model="hog"):
    if model == "cnn":
        return cnn_face_detector(img, number_of_times_to_upsample)
    else:
        return face_detector(img, number_of_times_to_upsample)


# step2 获取人脸68特征点
def _raw_face_landmarks(face_image, faces_locations=None):
    if faces_locations is None:
        faces_locations = _raw_face_locations(face_image)
    pose_predictor = pose_predictor_68_point
    return [pose_predictor(face_image, face_location) for face_location in faces_locations]


# step4 生成128维向量编码
def face_encodings(face_image, known_face_locations=None, num_jitters=3):
    raw_landmarks = _raw_face_landmarks(face_image, known_face_locations)
    return [numpy.array(face_encoder.compute_face_descriptor(face_image, raw_landmark_set, num_jitters)) for
            raw_landmark_set in raw_landmarks]


file_name = "image/001.jpg"
im = PIL.Image.open(file_name)
image = numpy.array(im)

encodings = face_encodings(image)
for code in encodings:
    print(str(code))
